import os
import csv
import re
import sys

def parse_output_file(filepath):
    """
    解析单个 output_batch*_threads*.txt 文件
    返回：列表，每个元素是 (op_name, p50, p99)
    """
    results = []
    with open(filepath, 'r') as f:
        lines = f.readlines()

    for line in lines:
        match = re.match(r'\|\s*(.*?)\s*\|\s*([\d\.]+)\s*ms\s*\|\s*([\d\.]+)\s*ms\s*\|', line)
        if match:
            name = match.group(1).strip()
            p50 = float(match.group(2))
            p99 = float(match.group(3))
            results.append((name, p50, p99))

    return results

def main(results_dir, output_csv):
    """
    遍历 results_raw 目录，生成总的 CSV
    """
    rows = []

    for device_type in os.listdir(results_dir):
        device_dir = os.path.join(results_dir, device_type)
        if not os.path.isdir(device_dir):
            continue

        for filename in os.listdir(device_dir):
            if not filename.startswith('output_batch'):
                continue

            filepath = os.path.join(device_dir, filename)

            # 提取 batch_size 和 threads_host
            match = re.match(r'output_batch(\d+)_threads(\d+)\.txt', filename)
            if not match:
                continue
            batch_size = int(match.group(1))
            threads_host = int(match.group(2))

            # 解析文件内容
            ops = parse_output_file(filepath)

            for op_name, p50, p99 in ops:
                rows.append({
                    'device_type': device_type,
                    'batch_size': batch_size,
                    'threads_host': threads_host,
                    'op_name': op_name,
                    'p50_ms': p50,
                    'p99_ms': p99
                })

    # 写入 CSV
    with open(output_csv, 'w', newline='') as csvfile:
        fieldnames = ['device_type', 'batch_size', 'threads_host', 'op_name', 'p50_ms', 'p99_ms']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        writer.writeheader()
        for row in rows:
            writer.writerow(row)

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python3 generate_summary_csv.py <results_dir> <output_csv>")
        sys.exit(1)

    results_dir = sys.argv[1]
    output_csv = sys.argv[2]

    main(results_dir, output_csv)
